{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "e:\\python27\\lib\\site-packages\\ipykernel\\__main__.py:10: FutureWarning: Categorical.from_array is deprecated, use Categorical instead\n"
     ]
    }
   ],
   "source": [
    "import pandas\n",
    "\n",
    "# Set index_col to False to avoid pandas thinking that the first column is row indexes (it's age).\n",
    "columns = [\"age\", \"workclass\", \"fnlwgt\", \"education\", \"education_num\", \"marital_status\", \"occupation\", \"relationship\", \"race\", \"sex\", \n",
    "           \"capital_gain\", \"capital_loss\", \"hours_per_week\", \"native_country\", \"high_income\"]\n",
    "#income = pandas.read_csv(\"income.csv\", index_col=False,names=columns)\n",
    "income = pandas.read_csv(\"income.csv\", names=columns)\n",
    "\n",
    "for name in [\"education\", \"marital_status\", \"occupation\", \"relationship\", \"race\", \"sex\", \"native_country\", \"high_income\",\"workclass\"]:\n",
    "    col = pandas.Categorical.from_array(income[name])\n",
    "    income[name] = col.codes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'train' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-2-343839d9c681>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0mclf\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mDecisionTreeClassifier\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mrandom_state\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mmin_samples_leaf\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;36m2\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m \u001b[0mclf\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mfit\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtrain\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mcolumns\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtrain\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;34m\"high_income\"\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      8\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[0mclf2\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mDecisionTreeClassifier\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mrandom_state\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mmax_depth\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;36m5\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'train' is not defined"
     ]
    }
   ],
   "source": [
    "from sklearn.tree import DecisionTreeClassifier\n",
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "columns = [\"age\", \"workclass\", \"education_num\", \"marital_status\", \"occupation\", \"relationship\", \"race\", \"sex\", \"hours_per_week\", \"native_country\"]\n",
    "\n",
    "clf = DecisionTreeClassifier(random_state=1, min_samples_leaf=2)\n",
    "clf.fit(train[columns], train[\"high_income\"])\n",
    "\n",
    "clf2 = DecisionTreeClassifier(random_state=1, max_depth=5)\n",
    "clf2.fit(train[columns], train[\"high_income\"])\n",
    "predictions = clf.predict(test[columns])\n",
    "print(roc_auc_score(test[\"high_income\"], predictions))\n",
    "\n",
    "predictions = clf2.predict(test[columns])\n",
    "print(roc_auc_score(test[\"high_income\"], predictions))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
